package com.zhuqian.myapplication

import android.app.Activity
import android.content.Context
import android.content.Intent
import android.graphics.*
import android.media.ExifInterface
import android.net.Uri
import android.os.Bundle
import android.provider.MediaStore
import android.widget.*
import androidx.activity.result.contract.ActivityResultContracts
import androidx.appcompat.app.AppCompatActivity
import androidx.core.content.FileProvider
import com.example.myapplication.R
import pl.droidsonroids.gif.GifDrawable
import pl.droidsonroids.gif.GifImageView
import java.io.File
import kotlin.math.min


class MainActivity : AppCompatActivity() {
    //    var mBitmap: Bitmap? = null
    var detector: Detector? = null
    var take_photo_btn: Button? = null
    var select_img_btn: Button? = null
    var detect_img_btn: Button? = null
    var resultTextView: TextView? = null
    var imageView: ScaledImageView? = null

    var is_detect = false


    fun show_msg(msg: String) {
        Toast.makeText(this, msg, Toast.LENGTH_SHORT).show()
    }

    fun drawRectOnBitmap(results: ArrayList<Result>, mBitmap: Bitmap) {
        val canvas = Canvas(mBitmap)
        val mPaintRectangle = Paint()
        mPaintRectangle.setColor(Color.GREEN)
        mPaintRectangle.strokeWidth = 5f
        mPaintRectangle.style = Paint.Style.FILL

        for ((classIndex, score, rect) in results) {
            val circle_x = (rect.left + rect.right) / 2
            val circle_y = (rect.top + rect.bottom) / 2
            val radius = min(rect.right - rect.left, rect.bottom - rect.top) * 1.0f / 2 * 0.8f
            canvas.drawCircle(circle_x.toFloat(), circle_y.toFloat(), radius, mPaintRectangle)
            canvas.drawRect(
                rect.left.toFloat(),
                rect.top.toFloat(),
                rect.left.toFloat(),
                rect.top.toFloat(),
                mPaintRectangle
            )
        }
    }


    private fun init_model() {

        detector = Detector(
            this,
            path = "dataset-v4-yolov8s-1024-v2.torchscript",
            img_size = 1024,
            type = "yolov8",
            class_num = 1,
            iou = 0.29f,
            conf = 0.29f,
            nms_limit = 9999,
        )
    }

    private fun init() {

        resultTextView = findViewById(R.id.resultTextView)
        select_img_btn = findViewById(R.id.select_img)
        take_photo_btn = findViewById(R.id.take_photo)
        detect_img_btn = findViewById(R.id.detect_img)
        init_take_photo_button()
        init_select_img_button()
        init_detect_img_button()
        Thread {
            init_model()
            ThreadUtils.runOnMainThread {
                ThreadUtils.runOnMainThread {
                    disable_all_button()
                }
                init_model() // 加载模型
                ThreadUtils.runOnMainThread {
                    init_bar()  // 模型初始化完毕后初始化参数调节条
                    enable_all_button()
                }
            }
        }.start()
    }

    fun set_text(text: String) {
        resultTextView?.visibility = TextView.VISIBLE
        resultTextView?.text = text
    }

    fun disable_all_button() {
        take_photo_btn?.isEnabled = false
        select_img_btn?.isEnabled = false
        detect_img_btn?.isEnabled = false
    }

    fun enable_all_button() {
        take_photo_btn?.isEnabled = true
        select_img_btn?.isEnabled = true
        detect_img_btn?.isEnabled = true
    }

    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        setContentView(R.layout.activity_main)
        imageView = findViewById(R.id.imageView)
        init()
    }

    private fun init_detect_img_button(){

        detect_img_btn?.setOnClickListener {
            if (is_detect) {
                return@setOnClickListener
            }
            imageView?.let {
                val bm = it.getBitmapFromRect()
                bm?.let {
                    detect_and_draw(it)
                }
            }
        }
    }

    private fun init_take_photo_button() {
        val outputImage = File(externalCacheDir, "output_image.jpg")
        if (outputImage.exists()) {
            outputImage.delete()
        }
        outputImage.createNewFile()
        val imageUri = FileProvider.getUriForFile(this, "com.example.myapplication.fileprovider", outputImage)
        val ar1 = registerForActivityResult(ActivityResultContracts.StartActivityForResult()) { result ->
            if (result.resultCode == Activity.RESULT_OK) {
                val bitmap = BitmapFactory.decodeStream(contentResolver.openInputStream(imageUri))
                val img_path = outputImage.absolutePath
                val bitmap2 = rotateIfRequired(bitmap, img_path)
                set_image(bitmap2)
            }
        }
        take_photo_btn!!.setOnClickListener {
            // 创建File对象，用于存储拍照后的图片
            val intent = Intent("android.media.action.IMAGE_CAPTURE")
            intent.putExtra(MediaStore.EXTRA_OUTPUT, imageUri)
            ar1.launch(intent)
        }
    }

    private fun init_select_img_button() {
        // 从相册选择图片
        val ar2 = registerForActivityResult(ActivityResultContracts.StartActivityForResult()) { result ->
            if (result.resultCode == Activity.RESULT_OK) {
                val uri = result.data?.data
                uri?.let {
                    val bitmap = getBitmapFromUri(it)
                    set_image(bitmap)
                }
            }
        }
        select_img_btn?.setOnClickListener {
            // 打开文件选择器
            val intent = Intent(Intent.ACTION_OPEN_DOCUMENT)
            intent.addCategory(Intent.CATEGORY_OPENABLE)
            // 指定只显示图片
            intent.type = "image/*"
            ar2.launch(intent)
        }
    }

    private fun rotateIfRequired(bitmap: Bitmap, img_path: String): Bitmap {
        val exif = ExifInterface(img_path)
        val orientation = exif.getAttributeInt(
            ExifInterface.TAG_ORIENTATION,
            ExifInterface.ORIENTATION_NORMAL
        )
        return when (orientation) {
            ExifInterface.ORIENTATION_ROTATE_90 -> rotateBitmap(bitmap, 90)
            ExifInterface.ORIENTATION_ROTATE_180 -> rotateBitmap(bitmap, 180)
            ExifInterface.ORIENTATION_ROTATE_270 -> rotateBitmap(bitmap, 270)
            else -> bitmap
        }
    }

    private fun rotateBitmap(bitmap: Bitmap, degree: Int): Bitmap {
        val matrix = Matrix()
        matrix.postRotate(degree.toFloat())
        val rotatedBitmap = Bitmap.createBitmap(
            bitmap, 0, 0, bitmap.width, bitmap.height,
            matrix, true
        )
        bitmap.recycle() // 将不再需要的Bitmap对象回收
        return rotatedBitmap
    }

    fun set_image(bitmap: Bitmap?) {
        resultTextView?.visibility = TextView.INVISIBLE
        imageView?.let {
            it.setBitmapOnView(bitmap)
            it.invalidate()
            is_detect = false
        }
    }

    fun detect_and_draw(bitmap: Bitmap) {
        val copy_bitmap = bitmap.copy(Bitmap.Config.ARGB_8888, true)
        if (detector != null) {
            disable_all_button()
            start_gif()
            // 模型预测
            Thread {
                val results = detector!!.detect_bitmap(copy_bitmap)
                is_detect = true
                // 绘制预测结果
                drawRectOnBitmap(results, copy_bitmap)
                ThreadUtils.runOnMainThread {
                    imageView?.setBitmapInRect(copy_bitmap)
                    set_text("检测数目: ${results.size}")
                    enable_all_button()
                    end_gif()
                }
            }.start()
        } else {
            show_msg("模型未加载")
        }
    }


    private fun init_bar() {
        val confBar: SeekBar = findViewById(R.id.confBar)
        val confTextView: TextView = findViewById(R.id.confTextView)
        val iouBar: SeekBar = findViewById(R.id.iouBar)
        val iouTextView: TextView = findViewById(R.id.iouTextView)
        val iou_cache: Float = get_cache_data("iou_cache", "0.29").toFloat()
        val conf_cache: Float = get_cache_data("conf_cache", "0.29").toFloat()
        detector?.let {
            it.conf = conf_cache
            it.iou = iou_cache
        }
        // 设置置信度参数调节条
        confTextView.text = String.format("%.2f", conf_cache)
        iouBar.progress = (iou_cache * 100).toInt()
        // 设置IOU参数调节条
        iouTextView.text = String.format("%.2f", iou_cache)
        confBar.progress = (conf_cache * 100).toInt()

        confBar.setOnSeekBarChangeListener(object : SeekBar.OnSeekBarChangeListener {
            override fun onProgressChanged(
                seekBar: SeekBar,
                progress: Int,
                b: Boolean
            ) {
                confTextView.text = String.format("%.2f", progress / 100.0f)
            }

            override fun onStartTrackingTouch(seekBar: SeekBar?) {

            }

            override fun onStopTrackingTouch(seekBar: SeekBar) {
                detector?.let {
                    it.conf = seekBar.progress / 100.0f
                }
                // 存储置信度参数
                set_cache_data("conf_cache", "${seekBar.progress / 100.0f}")
            }
        })


        iouBar.setOnSeekBarChangeListener(object : SeekBar.OnSeekBarChangeListener {
            override fun onProgressChanged(
                seekBar: SeekBar,
                progress: Int,
                b: Boolean
            ) {
                iouTextView.text = String.format("%.2f", progress / 100.0f)
            }

            override fun onStartTrackingTouch(seekBar: SeekBar?) {

            }

            override fun onStopTrackingTouch(seekBar: SeekBar) {
                detector?.let {
                    it.iou = seekBar.progress / 100.0f
                }
                // 存储IOU参数
                set_cache_data("iou_cache", "${seekBar.progress / 100.0f}")
            }
        })

    }


    private fun get_cache_data(key: String, default: String): String {
        val sp = getSharedPreferences("data", Context.MODE_PRIVATE)
        return sp.getString(key, default)!!
    }

    private fun set_cache_data(key: String, value: String) {
        val sp = getSharedPreferences("data", Context.MODE_PRIVATE)
        val editor = sp.edit()
        editor.putString(key, value)
        editor.apply()
    }

    var gifDrawable: GifDrawable? = null
    var gifImageView: GifImageView? = null

    private fun start_gif() {
        if (gifDrawable == null) {
            gifImageView = findViewById<GifImageView>(R.id.loading)
            gifDrawable = GifDrawable(resources, R.drawable.loading)
            gifImageView?.setImageDrawable(gifDrawable)
        }
        gifImageView?.visibility = ImageView.VISIBLE
        gifDrawable?.start()

    }

    private fun end_gif() {
        gifImageView?.visibility = ImageView.INVISIBLE
        gifDrawable?.stop()
    }

    private fun getBitmapFromUri(uri: Uri) = contentResolver
        .openFileDescriptor(uri, "r")?.use {
            BitmapFactory.decodeFileDescriptor(it.fileDescriptor)
        }
}